package com.niodata.dt.fs;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * an output stream byte buffer cache,convert output to inputstream.
 */
public class OutToInputBufferedStream extends OutputStream {

  private final InputStream inputStream = new ByteBufferInputStream();
  private final LinkedBlockingQueue<ByteBuffer> byteBuffers;
  private volatile boolean writeClosed;
  private volatile boolean readAll;
  private final int bufferSize;
  private ByteBuffer writeBuffer;
  private ByteBuffer readBuffer;
  private final int maxBuffers = 4;
  private boolean waitUntilAllRead = false;

  public OutToInputBufferedStream() {
    this(8 * 1024 * 1024, false);
  }

  public OutToInputBufferedStream(boolean waitUntilAllRead) {
    this(8 * 1024 * 1024, waitUntilAllRead);
  }

  public OutToInputBufferedStream(int cacheSize) {
    this(cacheSize, false);
  }

  public OutToInputBufferedStream(int cacheSize, boolean waitUntilAllRead) {
    this.bufferSize = cacheSize;
    byteBuffers = new LinkedBlockingQueue<>(maxBuffers);
    this.waitUntilAllRead = waitUntilAllRead;
  }

  @Override
  public void write(byte[] bytes) throws IOException {
    //bytes written
    int len = 0;
    while (len < bytes.length) {
      if (writeBuffer == null) {
        writeBuffer = ByteBuffer.allocate(bufferSize);
      }
      if (bytes.length <= writeBuffer.remaining()) {
        writeBuffer.put(bytes);
        len = len + bytes.length;
      } else {
        int remaining = writeBuffer.remaining();
        writeBuffer.put(bytes, len, remaining);
        len = len + remaining;
      }
      if (writeBuffer.position() == writeBuffer.capacity()) {
        writeBuffer.flip();
        try {
          byteBuffers.put(writeBuffer);
        } catch (InterruptedException e) {
          throw new RuntimeException(e);
        }
        writeBuffer = null;
      }
    }
  }

  @Override
  public void write(int b) throws IOException {
    if (writeBuffer == null) {
      writeBuffer = ByteBuffer.allocate(bufferSize);
    }
    writeBuffer.put((byte) b);
    if (writeBuffer.position() == writeBuffer.capacity()) {
      writeBuffer.flip();
      try {
        byteBuffers.put(writeBuffer);
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
      writeBuffer = ByteBuffer.allocate(bufferSize);
    }
  }

  public InputStream getInputStream() {
    return inputStream;
  }

  @Override
  public void close() throws IOException {
    if (writeBuffer != null) {
      writeBuffer.flip();
      try {
        byteBuffers.put(writeBuffer);
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
      writeBuffer = null;
    }
    writeClosed = true;
    long maxWait = 60 * 1000;
    while (!readAll) {
      try {
        Thread.sleep(1);
        maxWait--;
        if (maxWait <= 0) {
          throw new RuntimeException("output wait error");
        }
      } catch (InterruptedException e) {
        throw new RuntimeException(e);
      }
    }
  }

  private class ByteBufferInputStream extends InputStream {

    @Override
    public int read() throws IOException {
      if (readBuffer != null && readBuffer.hasRemaining()) {
        int b = readBuffer.get();
        if (!readBuffer.hasRemaining()) {
          readBuffer = null;
        }
        if (b < 0) {
          b = b + 256;
        }
        return b;
      }
      while (readBuffer == null) {
        readBuffer = byteBuffers.poll();
        if (readBuffer == null && writeClosed) {
          readAll = true;
          return -1;
        }
        try {
          Thread.sleep(2);
        } catch (InterruptedException ex) {
          throw new RuntimeException(ex);
        }
      }
      int b =  readBuffer.get();
      if (b < 0) {
        b = b + 256;
      }
      return b;
    }
  }
}
